#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
HiAgent Chat SDK 数据模型模块

定义了SDK中使用的各种数据模型类，用于结构化API响应数据。
"""

from dataclasses import dataclass, field
from typing import Dict, List, Optional, Any, Union
from datetime import datetime


@dataclass
class Conversation:
    """对话会话模型"""
    app_conversation_id: str
    conversation_name: str
    create_time: Optional[str] = None
    last_chat_time: Optional[str] = None
    empty_conversation: bool = False
    is_top: bool = False
    
    @classmethod
    def from_dict(cls, data: dict) -> 'Conversation':
        """从字典创建Conversation对象"""
        return cls(
            app_conversation_id=data.get('AppConversationID', ''),
            conversation_name=data.get('ConversationName', ''),
            create_time=data.get('CreateTime'),
            last_chat_time=data.get('LastChatTime'),
            empty_conversation=data.get('EmptyConversation', False),
            is_top=data.get('IsPinned', False)
        )


@dataclass
class FileInfo:
    """文件信息模型"""
    name: str
    path: str
    size: int
    url: str
    
    @classmethod
    def from_dict(cls, data: dict) -> 'FileInfo':
        """从字典创建FileInfo对象"""
        return cls(
            name=data.get('Name', ''),
            path=data.get('Path', ''),
            size=data.get('Size', 0),
            url=data.get('Url', '')
        )


@dataclass
class QueryExtends:
    """查询扩展信息模型"""
    files: List[FileInfo] = field(default_factory=list)
    
    def to_dict(self) -> dict:
        """转换为字典格式"""
        return {
            'Files': [
                {
                    'Name': f.name,
                    'Path': f.path,
                    'Size': f.size,
                    'Url': f.url
                } for f in self.files
            ]
        }


@dataclass
class Message:
    """消息模型"""
    message_id: str
    conversation_id: str
    content: str
    role: str  # 'user' or 'assistant'
    create_time: Optional[str] = None
    feedback: Optional[str] = None  # 'like' or 'dislike'
    message_type: Optional[str] = None
    status: Optional[str] = None
    retriever_resource: bool = False  # 是否使用了知识库资源
    knowledge_references: List[Dict[str, Any]] = field(default_factory=list)  # 知识库引用信息
    think_messages: List[Dict[str, Any]] = field(default_factory=list)  # 思考过程信息
    
    def _extract_knowledge_references(self, tracing_data) -> List[Dict[str, Any]]:
        """从TracingJsonStr中提取知识库引用信息

        Args:
            tracing_data: 可能是字符串(JSON)、列表或None
        """
        if not tracing_data:
            return []

        try:
            import json
            print(f"[DEBUG] 开始解析TracingJsonStr, 类型: {type(tracing_data)}")

            # 处理不同的数据类型
            if isinstance(tracing_data, str):
                # 如果是字符串，尝试解析JSON
                events = json.loads(tracing_data)
            elif isinstance(tracing_data, list):
                # 如果已经是列表，直接使用
                events = tracing_data
            else:
                print(f"[DEBUG] 不支持的数据类型: {type(tracing_data)}")
                return []

            print(f"[DEBUG] 事件数量: {len(events)}")
            references = []

            for event in events:
                event_type = event.get('event')
                print(f"[DEBUG] 处理事件: {event_type}")

                # 检查agent_thought_end事件中的observation字段
                if event_type == 'agent_thought_end':
                    observation = event.get('observation')
                    if observation:
                        try:
                            # observation是JSON字符串，需要解析
                            obs_data = json.loads(observation)
                            docs = obs_data.get('outputList', [])
                            print(f"[DEBUG] 从observation找到文档数量: {len(docs)}")

                            for i, doc in enumerate(docs, 1):
                                metadata = doc.get('metadata', {})

                                reference = {
                                    'reference_id': i,
                                    'slice_id': doc.get('slice_id', i),
                                    'document_name': metadata.get('document_name', '未知文档'),
                                    'similarity_score': metadata.get('score', 0),
                                    'content_preview': doc.get('output', '')[:200] + '...' if doc.get('output') else '',
                                    'document_url': metadata.get('document_obs_url', ''),
                                    'page_number': None,
                                    'full_content': doc.get('output', '')
                                }

                                # 提取页码信息
                                orientations = metadata.get('orientations', [])
                                if orientations:
                                    reference['page_number'] = orientations[0].get('page_number')

                                references.append(reference)
                        except json.JSONDecodeError:
                            print(f"[DEBUG] 无法解析observation JSON: {observation[:100]}...")

                # 兼容旧格式：直接的knowledge_retrieve_end事件
                elif event_type == 'knowledge_retrieve_end':
                    docs = event.get('docs', {}).get('outputList', [])
                    print(f"[DEBUG] 从knowledge_retrieve_end找到文档数量: {len(docs)}")

                    for i, doc in enumerate(docs, 1):
                        metadata = doc.get('metadata', {})

                        reference = {
                            'reference_id': i,
                            'slice_id': doc.get('slice_id', i),
                            'document_name': metadata.get('document_name', '未知文档'),
                            'similarity_score': metadata.get('score', 0),
                            'content_preview': doc.get('output', '')[:200] + '...' if doc.get('output') else '',
                            'document_url': metadata.get('document_obs_url', ''),
                            'page_number': None,
                            'full_content': doc.get('output', '')
                        }

                        # 提取页码信息
                        orientations = metadata.get('orientations', [])
                        if orientations:
                            reference['page_number'] = orientations[0].get('page_number')

                        references.append(reference)

            return references
        except Exception:
            return []

    def _extract_think_messages(self, tracing_data) -> List[Dict[str, Any]]:
        """从TracingJsonStr中提取思考过程信息

        Args:
            tracing_data: 可能是字符串(JSON)、列表或None

        Returns:
            List[Dict]: 思考过程信息列表
        """
        if not tracing_data:
            return []

        try:
            import json
            print(f"[DEBUG] 开始提取think_messages, 类型: {type(tracing_data)}")

            # 处理不同的数据类型
            if isinstance(tracing_data, str):
                events = json.loads(tracing_data)
            elif isinstance(tracing_data, list):
                events = tracing_data
            else:
                print(f"[DEBUG] 不支持的数据类型: {type(tracing_data)}")
                return []

            think_messages = []

            for event in events:
                event_type = event.get('event')
                if event_type == 'think_message':
                    think_message = {
                        'id': event.get('id', ''),
                        'content': event.get('answer', ''),
                        'created_at': event.get('created_at', 0),
                        'conversation_id': event.get('conversation_id', ''),
                        'task_id': event.get('task_id', ''),
                        'event_type': 'think_message'
                    }
                    think_messages.append(think_message)
                    print(f"[DEBUG] 添加思考消息: {len(think_message.get('content', ''))}字符")

            print(f"[DEBUG] 提取到{len(think_messages)}条think_messages")
            return think_messages
        except Exception as e:
            print(f"[DEBUG] 提取think_messages错误: {e}")
            return []

    @classmethod
    def from_dict(cls, data: dict) -> 'Message':
        """从字典创建Message对象，支持多种API格式"""
        # 标准格式（直接映射）
        if 'MessageID' in data or 'Content' in data:
            return cls(
                message_id=data.get('MessageID', ''),
                conversation_id=data.get('ConversationID', ''),
                content=data.get('Content', ''),
                role=data.get('Role', ''),
                create_time=data.get('CreateTime'),
                feedback=data.get('Feedback'),
                message_type=data.get('MessageType'),
                status=data.get('Status')
            )

        # HiAgent格式（需要从嵌套结构提取）
        elif 'AnswerInfo' in data and 'Query' in data:
            answer_info = data.get('AnswerInfo', {})
            created_timestamp = answer_info.get('CreatedTime', 0)

            # 转换时间戳为字符串
            create_time = None
            if created_timestamp:
                try:
                    from datetime import datetime
                    create_time = datetime.fromtimestamp(created_timestamp).strftime('%Y-%m-%d %H:%M:%S')
                except:
                    create_time = str(created_timestamp)

            # 创建消息对象
            message = cls(
                message_id=answer_info.get('MessageID', ''),
                conversation_id=data.get('ConversationID', ''),
                content=answer_info.get('Answer', ''),
                role='assistant',
                create_time=create_time,
                feedback='like' if answer_info.get('Like', 0) > 0 else None,
                message_type='text',
                status='completed',
                retriever_resource=answer_info.get('RetrieverResource', False)
            )

            # 提取知识库引用信息和思考过程
            tracing_data = answer_info.get('TracingJsonStr')
            print(f"[DEBUG] TracingJsonStr 类型: {type(tracing_data)}, 内容: {tracing_data}")
            if tracing_data:
                references = message._extract_knowledge_references(tracing_data)
                think_messages = message._extract_think_messages(tracing_data)
                print(f"[DEBUG] 提取到的引用数量: {len(references)}")
                print(f"[DEBUG] 提取到的思考消息数量: {len(think_messages)}")
                message.knowledge_references = references
                message.think_messages = think_messages

            return message

        # 兜底：返回空Message
        else:
            return cls(
                message_id='',
                conversation_id=data.get('ConversationID', ''),
                content='',
                role='',
                create_time=None,
                feedback=None,
                message_type=None,
                status=None
            )


@dataclass
class ChatResponse:
    """聊天响应模型"""
    message_id: str
    content: str
    conversation_id: str
    status: str = 'completed'
    error: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)
    knowledge_references: List[Dict[str, Any]] = field(default_factory=list)
    think_messages: List[Dict[str, Any]] = field(default_factory=list)
    
    @classmethod
    def from_dict(cls, data: dict) -> 'ChatResponse':
        """从字典创建ChatResponse对象"""
        # 支持多种API响应格式
        message_id = data.get('MessageID') or data.get('id') or data.get('task_id', '')
        content = data.get('Content') or data.get('answer', '')
        conversation_id = data.get('ConversationID') or data.get('conversation_id', '')
        status = data.get('Status') or data.get('event', 'completed')

        # 创建基本的ChatResponse对象
        response = cls(
            message_id=message_id,
            content=content,
            conversation_id=conversation_id,
            status=status,
            error=data.get('Error'),
            metadata=data.get('Metadata', {})
        )

        # 处理think_messages（直接从数据中获取）
        think_messages = data.get('think_messages', [])
        if think_messages:
            # 如果是字符串列表，转换为字典格式
            formatted_think_messages = []
            for i, msg in enumerate(think_messages):
                if isinstance(msg, str):
                    formatted_think_messages.append({
                        'id': f'think_{i}',
                        'content': msg,
                        'created_at': None,
                        'conversation_id': conversation_id,
                        'task_id': message_id,
                        'event_type': 'think_message'
                    })
                else:
                    formatted_think_messages.append(msg)
            response.think_messages = formatted_think_messages

        return response


@dataclass
class WorkflowResult:
    """工作流执行结果模型"""
    task_id: str
    status: str
    result: Optional[Dict[str, Any]] = None
    error: Optional[str] = None
    progress: float = 0.0
    
    @classmethod
    def from_dict(cls, data: dict) -> 'WorkflowResult':
        """从字典创建WorkflowResult对象"""
        return cls(
            task_id=data.get('TaskID', ''),
            status=data.get('Status', ''),
            result=data.get('Result'),
            error=data.get('Error'),
            progress=data.get('Progress', 0.0)
        )


@dataclass
class LongMemory:
    """长期记忆模型"""
    memory_id: str
    content: str
    create_time: Optional[str] = None
    update_time: Optional[str] = None
    
    @classmethod
    def from_dict(cls, data: dict) -> 'LongMemory':
        """从字典创建LongMemory对象"""
        return cls(
            memory_id=data.get('MemoryID', ''),
            content=data.get('Content', ''),
            create_time=data.get('CreateTime'),
            update_time=data.get('UpdateTime')
        )


@dataclass
class UserVariable:
    """用户变量模型"""
    key: str
    value: str
    description: Optional[str] = None
    
    @classmethod
    def from_dict(cls, data: dict) -> 'UserVariable':
        """从字典创建UserVariable对象"""
        return cls(
            key=data.get('Key', ''),
            value=data.get('Value', ''),
            description=data.get('Description')
        )